package tree;

import java.util.*;

public class HaffmanTree {

    static class HaffmanTreeNode extends Node {

        Character c;

        public HaffmanTreeNode(int value, Character c) {
            super(value);
            this.c = c;
        }
    }

    HaffmanTreeNode root;

    Map<Character, String> incodeMap = new HashMap<>(); // 编码表

    Map<String, Character> decodeMap = new HashMap<>(); // 译码表

    /**
     * 根据字符串构建哈夫曼树
     * @param s 字符串内容
     */
    public void buildTree(String s) {
        // 如果只有一个值直接返回
        if (s.length() == 1) {
            this.root = new HaffmanTreeNode(1, s.charAt(0));
        }
        // 统计每个字符出现次数
        Map<Character, Integer> charMap = new HashMap<>();
        for (char c : s.toCharArray()) {
            charMap.merge(c, 1, Integer::sum);
        }
        // 将列表中的数字构建节点存放到集合中
        List<HaffmanTreeNode> nodeList = new LinkedList<>();
        charMap.forEach((c,num) -> nodeList.add(new HaffmanTreeNode(num, c)));

        // 对集合中的节点进行排序(冒泡排序)
        boolean flag = true;
        for (int i = nodeList.size() - 1;  i >= 0 && flag; --i) {
            flag = false;
            for (int j = 0; j < i; ++j) {
                if (nodeList.get(j).value > nodeList.get(j + 1).value) {
                    int tmp = nodeList.get(j).value;
                    nodeList.get(j).value = nodeList.get(j + 1).value;
                    nodeList.get(j + 1).value = tmp;
                    char tmpc = nodeList.get(j).c;
                    nodeList.get(j).c = nodeList.get(j + 1).c;
                    nodeList.get(j + 1).c = tmpc;
                    flag = true;
                }
            }
        }
        while (nodeList.size() > 1) {
            // 选出最小的两个
            HaffmanTreeNode node1 = nodeList.remove(0);
            HaffmanTreeNode node2 = nodeList.remove(0);
            // 构建新的节点
            HaffmanTreeNode node3 = new HaffmanTreeNode(node1.value + node2.value, null);
            node3.left = node1;
            node3.right = node2;
            // 把节点放入集合
            int idx = 0;
            for (; idx < nodeList.size(); idx++) {
                if (nodeList.get(idx).value > node3.value) {
                    break;
                }
            }
            nodeList.add(idx, node3);
        }
        // 返回最后剩下的节点
        this.root = nodeList.get(0);
        // 构建编码表和译码表
        getCodeMap(root, incodeMap, "");
        incodeMap.forEach((k, v) -> decodeMap.put(v, k));
    }

    /**
     * 获取编码表的辅助函数
     * @param node  树节点
     * @param charMap   编码表
     * @param origin    初始编码值
     */
    private void getCodeMap(HaffmanTreeNode node, Map<Character, String> charMap, String origin) {
        // 只保存叶子节点的编码值
        if (node.left == null && node.right == null) {
            charMap.put(node.c, origin);
        }
        if (node.left != null) {
            getCodeMap((HaffmanTreeNode) node.left, charMap, origin + "0");
        }
        if (node.right != null) {
            getCodeMap((HaffmanTreeNode) node.right, charMap, origin + "1");
        }
    }

    /**
     * 编码
     * @param input 输入内容
     * @return  编码后内容
     */
    public String incode(String input) {
        // 逐字符编码
        char[] chars = input.toCharArray();
        StringBuilder res = new StringBuilder();
        for (char c : chars) {
            res.append(this.incodeMap.get(c));
        }
        // 返回结果
        return res.toString();
    }

    /**
     * 解码
     * @param input 原内容
     * @return  解码后内容
     */
    public String decode(String input) {
        StringBuilder code = new StringBuilder();
        int idx = 0, maxLen = this.decodeMap.size() - 1;
        while (idx < input.length()) {
            // 最长匹配原则
            for (int i = maxLen; i > 0; --i) {
                if (idx + i <= input.length()) {
                    String substring = input.substring(idx, idx + i);
                    if (decodeMap.containsKey(substring)) {
                        code.append(decodeMap.get(substring));
                        idx += i;
                        break;
                    }
                }
            }
        }
        return code.toString();
    }

    public static void main(String[] args) {
        String input = "Hello World";
        HaffmanTree haffmanTree = new HaffmanTree();
        haffmanTree.buildTree(input);
        System.out.println(haffmanTree.incode(input));
        System.out.println(haffmanTree.decode(haffmanTree.incode(input)));
        /*
          10
         /  \
        4    6
            / \
           3   3
              / \
             1   2
         */
//        prePrint(haffmanTree, res);
//        for (int value : res) {
//            System.out.println(value + " ");
//        }
    }
}
